import gradio as gr
from sklearn.cluster import KMeans
from PIL import Image
import numpy as np

def extract_colors(image_path, color_count=5):
    # 打开图像
    image = Image.open(image_path)
    # 将图像转换为RGB模式
    image = image.convert('RGB')
    # 将图像数据转换为numpy数组
    data = np.array(image)
    # 将数据重塑为二维数组，每行代表一个像素，列代表RGB值
    data = data.reshape((data.shape[0] * data.shape[1], 3))
    
    # 使用KMeans算法提取颜色
    kmeans = KMeans(n_clusters=color_count)
    kmeans.fit(data)
    # 获取聚类中心，即主色调
    palette = kmeans.cluster_centers_
    
    # 将颜色转换为十六进制表示形式
    hex_palette = ["#{:02x}{:02x}{:02x}".format(*color.astype(int)) for color in palette]
    
    # 创建调色板图像，背景为白色
    palette_width = 250  # 调色板的总宽度
    palette_height = 50  # 调色板的高度
    palette_image = Image.new('RGB', (palette_width, palette_height), 'white')
    
    # 计算每个颜色块的宽度
    color_width = palette_width // color_count
    
    # 在调色板图像上绘制每个颜色块
    for i, color in enumerate(palette):
        # 计算颜色块的左上角坐标
        left = i * color_width
        top = 0
        # 绘制颜色块
        palette_image.paste(Image.new('RGB', (color_width, palette_height), tuple(color.astype(int))), (left, top))
    
    return palette_image, "\n".join(hex_palette)

# 创建Gradio界面
iface = gr.Interface(
    fn=extract_colors,
    inputs=gr.Image(type="filepath", label="上传图片"),
    outputs=[gr.Image(label="调色板"), gr.Textbox(label="色彩十六进制码")],
    title="KMeans调色板提取器",
    description="上传图像获取其调色板。",
    theme="default"
)

# 启动应用
iface.launch()